package org.dromara.common.ratelimiter.aspectj

import io.github.oshai.kotlinlogging.KotlinLogging
import org.aspectj.lang.JoinPoint
import org.aspectj.lang.annotation.Aspect
import org.aspectj.lang.annotation.Before
import org.aspectj.lang.reflect.MethodSignature
import org.dromara.common.core.constant.GlobalConstants
import org.dromara.common.core.exception.ServiceException
import org.dromara.common.core.utils.MessageUtils.message
import org.dromara.common.core.utils.ServletUtils.getClientIP
import org.dromara.common.core.utils.ServletUtils.getRequest
import org.dromara.common.core.utils.SpringUtils.getBeanFactory
import org.dromara.common.core.utils.StringUtils
import org.dromara.common.ratelimiter.annotation.RateLimiter
import org.dromara.common.ratelimiter.enums.LimitType
import org.dromara.common.redis.utils.RedisUtils.getClient
import org.dromara.common.redis.utils.RedisUtils.rateLimiter
import org.redisson.api.RateType
import org.springframework.context.expression.BeanFactoryResolver
import org.springframework.context.expression.MethodBasedEvaluationContext
import org.springframework.core.DefaultParameterNameDiscoverer
import org.springframework.core.ParameterNameDiscoverer
import org.springframework.expression.ExpressionParser
import org.springframework.expression.ParserContext
import org.springframework.expression.common.TemplateParserContext
import org.springframework.expression.spel.standard.SpelExpressionParser

/**
 * 限流处理
 *
 * @author Lion Li
 * @updater LikeYouDo
 * @date 2025/2/1 20:06
 */
@Aspect
class RateLimiterAspect {
    companion object {
        private val log = KotlinLogging.logger { }
    }

    /**
     * 定义spel表达式解析器
     */
    private val parser: ExpressionParser = SpelExpressionParser()

    /**
     * 定义spel解析模版
     */
    private val parserContext: ParserContext = TemplateParserContext()

    /**
     * 方法参数解析器
     */
    private val pnd: ParameterNameDiscoverer = DefaultParameterNameDiscoverer()

    @Before("@annotation(rateLimiter)")
    fun doBefore(point: JoinPoint, rateLimiter: RateLimiter) {
        val time = rateLimiter.time
        val count = rateLimiter.count
        val timeout = rateLimiter.timeout
        try {
            val combineKey = getCombineKey(rateLimiter, point)
            var rateType = RateType.OVERALL
            if (rateLimiter.limitType == LimitType.CLUSTER) {
                rateType = RateType.PER_CLIENT
            }
            val number = rateLimiter(combineKey, rateType, count, time, timeout)
            if (number == -1L) {
                var message = rateLimiter.message
                if (StringUtils.startsWith(message, "{") && StringUtils.endsWith(message, "}")) {
                    message = message(StringUtils.substring(message, 1, message.length - 1))
                }
                throw ServiceException(message)
            }
            log.info { "限制令牌 => ${count}, 剩余令牌 => ${number}, 缓存key => '${combineKey}'" }
        } catch (e: Exception) {
            if (e is ServiceException) {
                throw e
            } else {
                throw RuntimeException("服务器限流异常，请稍候再试", e)
            }
        }
    }

    private fun getCombineKey(rateLimiter: RateLimiter, point: JoinPoint): String {
        var key = rateLimiter.key
        // 判断 key 不为空 和 不是表达式
        if (StringUtils.isNotBlank(key) && StringUtils.containsAny(key, "#")) {
            val signature = point.signature as MethodSignature
            val targetMethod = signature.method
            val args = point.args
            val targetObj = point.target
            val context =
                MethodBasedEvaluationContext(targetObj, targetMethod, args, pnd)
            context.setBeanResolver(BeanFactoryResolver(getBeanFactory()))
            val expression =
                if (StringUtils.startsWith(key, parserContext.expressionPrefix)
                    && StringUtils.endsWith(key, parserContext.expressionSuffix)
                ) {
                    parser.parseExpression(key, parserContext)
                } else {
                    parser.parseExpression(key)
                }
            key = expression.getValue(context, String::class.java).toString()
        }
        val stringBuffer = StringBuilder(GlobalConstants.RATE_LIMIT_KEY)
        stringBuffer.append(getRequest().requestURI).append(":")
        if (rateLimiter.limitType == LimitType.IP) {
            // 获取请求ip
            stringBuffer.append(getClientIP()).append(":")
        } else if (rateLimiter.limitType == LimitType.CLUSTER) {
            // 获取客户端实例id
            stringBuffer.append(getClient().id).append(":")
        }
        return stringBuffer.append(key).toString()
    }
}
